/**
 * Copyright(c) 2016-2018 by Euler
 * All Rights Reserved
 */
package com.booter.gateway.predicate;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.handler.AsyncPredicate;
import org.springframework.cloud.gateway.handler.predicate.AbstractRoutePredicateFactory;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Base64;
import java.util.List;
import java.util.function.Predicate;

/**
 * @author lichunwen
 * 解决 requestbody 被截断的问题
 */
@Component("CustomReadBody")
@Slf4j
public class CustomReadBodyPredicateFactory
        extends AbstractRoutePredicateFactory<CustomReadBodyConfig> {

    private static final String TEST_ATTRIBUTE = "read_body_predicate_test_attribute";

    public static final String CACHE_OBJECT = "cachedRequestBodyObject";

    private static final List<HttpMessageReader<?>> messageReaders = HandlerStrategies
            .withDefaults().messageReaders();

    private CustomPredicate predicate = new CustomPredicate();


    public CustomReadBodyPredicateFactory() {
        super(CustomReadBodyConfig.class);
    }

    @Override
    public AsyncPredicate<ServerWebExchange> applyAsync(CustomReadBodyConfig config) {
        return exchange -> {
            Class<String> inClass = String.class;
            //检查缓存中有没有
            String cachedBody = exchange.getAttribute(CACHE_OBJECT);
            if (cachedBody != null) {
                try {
                    boolean test = this.predicate.test(cachedBody);
                    if (cachedBody != null) {
                        String str = new String(Base64.getDecoder().decode(cachedBody));
                        exchange.getAttributes().put(TEST_ATTRIBUTE, str);
                    }
                    return Mono.just(test);
                } catch (ClassCastException e) {
                    if (log.isDebugEnabled()) {
                        log.debug("Predicate test failed because class in predicate does not match the cached body object", e);
                    }
                }
                return Mono.just(false);
            }
            //缓存中没有时从request中获取
            return ServerWebExchangeUtils.cacheRequestBodyAndRequest(exchange, (serverHttpRequest) -> ServerRequest
                    .create(exchange.mutate().request(serverHttpRequest).build(), messageReaders)
                    .bodyToMono(inClass)
                    .doOnNext(s -> {
                        String str = new String(Base64.getDecoder().decode(s));
                        exchange.getAttributes().put(CACHE_OBJECT, str);
                    })
                    .map(objectValue -> {
                        String str = new String(Base64.getDecoder().decode(objectValue));
                        return this.predicate.test(str);
                    }));

        };
    }

    @Override
    public Predicate<ServerWebExchange> apply(CustomReadBodyConfig config) {
        throw new UnsupportedOperationException("CustomReadBodyPredicateFactory is only async.");
    }

}
